[None][refactor] Unify sparse attention framework with clean backend interfaces#12733
Open
lfr-0531 wants to merge 11 commits into
Open
[None][refactor] Unify sparse attention framework with clean backend interfaces#12733lfr-0531 wants to merge 11 commits into
lfr-0531 wants to merge 11 commits into
Conversation
lfr-0531
force-pushed
the
user/fanrongl/sparse-attention-refactor
branch
8 times, most recently
from
April 3, 2026 11:18
a1e4402 to
eaab4c3
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41657 [ run ] triggered by Bot. Commit: |
1 task
Collaborator
|
PR_Github #41657 [ run ] completed with state
|
lfr-0531
force-pushed
the
user/fanrongl/sparse-attention-refactor
branch
2 times, most recently
from
April 4, 2026 06:23
6d89705 to
4004b1a
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41813 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #41813 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41921 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #41921 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41962 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #41962 [ run ] completed with state
|
heyuhhh
reviewed
Apr 7, 2026
|
|
||
| TLLM_CHECK(host_kv_cache_pool_mapping.has_value()); | ||
| int32_t const layer_num = host_kv_cache_pool_mapping.value().size(0); | ||
|
|
Collaborator
There was a problem hiding this comment.
We can add an assert sentence here to make sure sparse_mla_topk_value is a valid value
heyuhhh
reviewed
Apr 7, 2026
| tensorrt_llm/_torch/attention_backend/sparse/dsa.py | | ||
| tensorrt_llm/_torch/attention_backend/sparse/kernel.py | | ||
| tensorrt_llm/_torch/attention_backend/sparse/rocket.py | | ||
| tensorrt_llm/_torch/attention_backend/sparse/utils.py | |
Collaborator
There was a problem hiding this comment.
Why we just remove there files but not add new files?
Collaborator
Author
There was a problem hiding this comment.
The removed files (dsa.py, kernel.py, rocket.py, utils.py) were in the legacy exclusion list because they had historical lint violations that were grandfathered in.
The new replacement files (dsa/backend.py, dsa/custom_ops.py, dsa/indexer.py, dsa/metadata.py, dsa/cache_manager.py, rocket/backend.py, skip_softmax/backend.py, params.py, etc.) are written from scratch and fully comply with current lint standards — ruff check passes on all of them without any exclusions. So they don't need to be added to the legacy list.
The same applies to legacy-files.txt, pyproject.toml, and ruff-legacy.toml — all four config files are auto-generated from legacy-files.txt via scripts/legacy_utils.py gen-configs.
lfr-0531
force-pushed
the
user/fanrongl/sparse-attention-refactor
branch
3 times, most recently
from
April 7, 2026 03:07
79c8913 to
53a6ee5
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61504 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61465 [ run ] completed with state |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py (1)
9-13: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the new rocketkv tests to the integration lists
test_triton_bmm,test_triton_rocket_paged_kt_cache_bmm, andtest_topk_kernelare new intests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py.- No matching entries appear under
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/, so these tests won’t be picked up by the integration harness.Coverage verdict: insufficient.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py` around lines 9 - 13, Add integration-list entries for test_triton_bmm, test_triton_rocket_paged_kt_cache_bmm, and test_topk_kernel from test_kernels.py in the appropriate test-db and qa lists, using the existing list format so the integration harness discovers all three tests.Source: Path instructions
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py (1)
727-728: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead no-op expression. Line 728 computes
num_kv_heads * num_heads_per_kvbut discards the result — it has no effect and looks like a leftover (possibly intended to assign anum_heads/total_headslocal). Remove it or assign it if the value is needed.♻️ Suggested cleanup
num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim = q.shape - num_kv_heads * num_heads_per_kv🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py` around lines 727 - 728, Remove the discarded num_kv_heads * num_heads_per_kv expression after the q.shape unpacking, unless a later computation requires that product; if needed, assign it to the appropriate existing heads-count local and use that value.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py`:
- Around line 9-13: Add integration-list entries for test_triton_bmm,
test_triton_rocket_paged_kt_cache_bmm, and test_topk_kernel from test_kernels.py
in the appropriate test-db and qa lists, using the existing list format so the
integration harness discovers all three tests.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py`:
- Around line 727-728: Remove the discarded num_kv_heads * num_heads_per_kv
expression after the q.shape unpacking, unless a later computation requires that
product; if needed, assign it to the appropriate existing heads-count local and
use that value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 81865ea1-94a0-455d-88fc-d90fc7d04e41
📒 Files selected for processing (75)
.pre-commit-config.yamldocs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.mddocs/source/developer-guide/sparse-attention-development-guide.mddocs/source/torch/adding_custom_kernels.mdlegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.pytensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.pytensorrt_llm/_torch/attention_backend/sparse/dsa/backend.pytensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.pytensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.pytensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.pytensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.pytensorrt_llm/_torch/attention_backend/sparse/dsa/module.pytensorrt_llm/_torch/attention_backend/sparse/dsa/params.pytensorrt_llm/_torch/attention_backend/sparse/hooks.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/params.pytensorrt_llm/_torch/attention_backend/sparse/registry.pytensorrt_llm/_torch/attention_backend/sparse/rocket.pytensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.pytensorrt_llm/_torch/attention_backend/sparse/rocket/backend.pytensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.pytensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.pytensorrt_llm/_torch/attention_backend/sparse/rocket/module.pytensorrt_llm/_torch/attention_backend/sparse/rocket/params.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_kernels.pytests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.pytests/unittest/_torch/attention/sparse/rocketkv/test_kernels.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/attention/sparse/test_sparse_attention.pytests/unittest/_torch/attention/sparse/test_sparse_mla_forward.pytests/unittest/_torch/attention/sparse/test_triton_topk.pytests/unittest/_torch/attention/test_attention_op_sync.pytests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.pytests/unittest/_torch/modeling/test_modeling_deepseekv4.pytests/unittest/_torch/modules/test_mla_helix.pytests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (2)
- tests/unittest/_torch/attention/sparse/test_triton_topk.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket.py
🚧 Files skipped from review as they are similar to previous changes (56)
- tensorrt_llm/_torch/attention_backend/sparse/registry.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/module.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/init.py
- tensorrt_llm/_torch/attention_backend/sparse/init.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/params.py
- tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/init.py
- tests/unittest/_torch/attention/sparse/dsa/test_kernels.py
- pyproject.toml
- tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/init.py
- legacy-files.txt
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py
- ruff-legacy.toml
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py
- tests/unittest/_torch/attention/test_attention_op_sync.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
- docs/source/torch/adding_custom_kernels.md
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py
- tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
- tests/unittest/_torch/modules/test_mla_helix.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py
- tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py
- tensorrt_llm/_torch/attention_backend/vanilla.py
- tensorrt_llm/_torch/models/modeling_deepseekv4.py
- docs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.md
- tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py
- tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
- tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.py
- tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
- docs/source/developer-guide/sparse-attention-development-guide.md
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/init.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py
- tensorrt_llm/_torch/attention_backend/sparse/hooks.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py
- tensorrt_llm/_torch/attention_backend/trtllm.py
- tensorrt_llm/_torch/attention_backend/fmha/fallback.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.py
- tensorrt_llm/_torch/attention_backend/sparse/params.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py
- tensorrt_llm/_torch/modules/attention.py
- tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
- tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py
- tensorrt_llm/_torch/modules/mla.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py (1)
727-728: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead expression statement.
num_kv_heads * num_heads_per_kvcomputes and discards a value with no effect. Remove it (or assign it if it was meant to be used).♻️ Proposed cleanup
num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim = q.shape - num_kv_heads * num_heads_per_kv🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py` around lines 727 - 728, Remove the unused expression statement `num_kv_heads * num_heads_per_kv` from the code immediately following the q.shape unpacking; leave the existing shape assignments unchanged unless the computed value is required by subsequent logic.tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py (1)
102-105: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing return-type annotations.
post_load_weights(Line 102),_qk_projection_and_rope(Line 107), andforward(Line 357) lack return-type annotations (-> None/-> torch.Tensor).As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions."✏️ Proposed fix
- def post_load_weights(self): + def post_load_weights(self) -> None: # V4 does not use the V3 fused fp32 wk+weights_proj GEMM, and the # base concat would now hit an fp32/bf16 dtype mismatch. return- def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): + def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:def forward( self, qr: torch.Tensor, hidden_states: torch.Tensor, metadata: DeepseekV4TrtllmAttentionMetadata, position_ids: torch.Tensor, pre_aux: Optional[ Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] ] = None, - ): + ) -> torch.Tensor:Also applies to: 107-115, 357-366
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py` around lines 102 - 105, Annotate post_load_weights with -> None, _qk_projection_and_rope with -> torch.Tensor, and forward with the appropriate torch.Tensor return annotation, preserving their existing behavior and signatures otherwise.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py`:
- Around line 208-224: The FP8_BLOCKWISE path in _update_k_cache_if_needed
currently invokes _update_k_cache even though that override is a no-op, leaving
the indexer K cache unpopulated. Implement the required FP8 blockwise cache
update so this path populates the cache, or remove the unreachable FP8 handling
including the k_scale assertion if FP8 blockwise is unsupported; keep the MXFP4
early return unchanged.
In `@tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py`:
- Around line 221-226: The test currently passes a flattened rank-2 view to
deepseek_v4_q_norm_fused_fp8 instead of exercising the documented 3D q_pe
contract. Update the invocation around q_pe and the custom op to pass the
original 3D tensor, or validate it through the sparse-MLA attention consumer,
while retaining assertions that verify the actual consumer path rather than only
the allocation.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py`:
- Around line 102-105: Annotate post_load_weights with -> None,
_qk_projection_and_rope with -> torch.Tensor, and forward with the appropriate
torch.Tensor return annotation, preserving their existing behavior and
signatures otherwise.
In `@tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py`:
- Around line 727-728: Remove the unused expression statement `num_kv_heads *
num_heads_per_kv` from the code immediately following the q.shape unpacking;
leave the existing shape assignments unchanged unless the computed value is
required by subsequent logic.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d21421a4-1c83-4620-bfaa-d0675f36744a
📒 Files selected for processing (75)
.pre-commit-config.yamldocs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.mddocs/source/developer-guide/sparse-attention-development-guide.mddocs/source/torch/adding_custom_kernels.mdlegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.pytensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.pytensorrt_llm/_torch/attention_backend/sparse/dsa/backend.pytensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.pytensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.pytensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.pytensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.pytensorrt_llm/_torch/attention_backend/sparse/dsa/module.pytensorrt_llm/_torch/attention_backend/sparse/dsa/params.pytensorrt_llm/_torch/attention_backend/sparse/hooks.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/params.pytensorrt_llm/_torch/attention_backend/sparse/registry.pytensorrt_llm/_torch/attention_backend/sparse/rocket.pytensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.pytensorrt_llm/_torch/attention_backend/sparse/rocket/backend.pytensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.pytensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.pytensorrt_llm/_torch/attention_backend/sparse/rocket/module.pytensorrt_llm/_torch/attention_backend/sparse/rocket/params.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_kernels.pytests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.pytests/unittest/_torch/attention/sparse/rocketkv/test_kernels.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/attention/sparse/test_sparse_attention.pytests/unittest/_torch/attention/sparse/test_sparse_mla_forward.pytests/unittest/_torch/attention/sparse/test_triton_topk.pytests/unittest/_torch/attention/test_attention_op_sync.pytests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.pytests/unittest/_torch/modeling/test_modeling_deepseekv4.pytests/unittest/_torch/modules/test_mla_helix.pytests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (2)
- tests/unittest/_torch/attention/sparse/test_triton_topk.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket.py
🚧 Files skipped from review as they are similar to previous changes (54)
- tensorrt_llm/_torch/attention_backend/sparse/rocket/init.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/module.py
- tensorrt_llm/_torch/attention_backend/sparse/registry.py
- ruff-legacy.toml
- docs/source/torch/adding_custom_kernels.md
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/init.py
- tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/params.py
- tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
- tensorrt_llm/_torch/attention_backend/sparse/init.py
- tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
- pyproject.toml
- tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py
- .pre-commit-config.yaml
- tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
- tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py
- tests/unittest/_torch/attention/test_attention_op_sync.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py
- tensorrt_llm/_torch/models/modeling_deepseekv4.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py
- tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
- tests/unittest/llmapi/test_llm_args.py
- tensorrt_llm/_torch/attention_backend/vanilla.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.py
- tensorrt_llm/_torch/attention_backend/sparse/hooks.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
- tensorrt_llm/llmapi/llm_args.py
- tests/unittest/_torch/modules/test_mla_helix.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
- tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py
- tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
- tensorrt_llm/_torch/attention_backend/trtllm.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py
- tensorrt_llm/_torch/attention_backend/fmha/fallback.py
- tensorrt_llm/_torch/attention_backend/interface.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py
- tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py
- tensorrt_llm/_torch/modules/attention.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/init.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py
- tests/unittest/_torch/attention/sparse/test_sparse_attention.py
- tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py
- tensorrt_llm/_torch/modules/mla.py
Collaborator
|
PR_Github #61533 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61504 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61541 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61533 [ run ] completed with state |
Shixiaowei02
approved these changes
Jul 24, 2026
Collaborator
|
PR_Github #61541 [ run ] completed with state
|
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Define validated sparse-attention hooks shared by Attention and MLA while keeping algorithm-specific module behavior under each sparse backend. Simplify DSA and DeepSeek-V4 module integration, move Rocket kernels into its backend directory, and align the related architecture tests. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Unify module-to-backend and backend-to-attention-op sparse runtime arguments. Split DeepSeek-V4 indexer, metadata, and parameter definitions, and route DSA prediction through the backend while preserving shared TopK buffer lifetime. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Allocate a single mixed-batch TopK buffer only when the model contains shared indexer layers. Reuse per-layer indexer routing to derive the metadata requirement and keep the buffer address stable across steps. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Keep split DSA indexer test paths under inference mode, matching the original integrated Indexer.forward contract. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Restore the shared lint configuration to match main and format the split sparse attention modules with Ruff. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61736 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61736 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61744 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61744 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dev Engineer Review
SparseBackendForwardArgsandSparseRuntimeParams, removingSparsePredictionand replacingSkipSoftmaxKernelParamsusage with runtime fields.tensorrt_llm/_torch/attention_backend/sparse/hooks.py, and rewired MLA/attention execution to route sparse execution through hooks and consume predicted indices fromforward_args.sparse_runtime_params.*.tensorrt_llm/_torch/attention_backend/sparse/<algorithm>/:sparse/dsa/package (backend, indexer, kernels, metadata, cache manager, custom CUDA-graph-friendly ops, params).sparse/deepseek_v4/package (runtime-only metadata, backend/indexer/kernels/compressor, params).sparse/rocket.pytosparse/rocket/package; the monolithicrocket.pywas deleted and replaced by package modules (backend/cache_manager/kernels/metadata/params).SparseRuntimeParamsacross multiple backends and FMHA fallbacks, including SkipSoftmax threshold scaling fields..pre-commit-config.yaml,legacy-files.txt,pyproject.toml,ruff-legacy.toml.tok < 0, page-range check), not necessarily on all negative/out-of-range intermediates—verify assumptions aboutblock_table/pool-page validity.on_update_kv_lens()and CUDA-graph capture paths via custom ops; ensure buffer object identity/stability and cache revalidation are consistent.sparse/rocket/package and that any stale references to the removedrocket.pyare gone.forward_args.sparse_runtime_params(not legacy prediction/kernel-param paths), especially where multiple FMHA/FlashInfer/MSA backends branch on “sparse vs dense plan” decisions.QA Engineer Review
No test changes (no
tests/,tests/integration/test_lists/,test-db/,qa/, orwaives.txtmodifications detected in the diff evidence available here).Description
TensorRT-LLM already provides a sparse-attention framework for DSA, DeepSeek-V4,
RocketKV, and SkipSoftmax, but algorithm-specific module, backend, prediction,
metadata, cache, and kernel logic had become mixed across the common Attention
and MLA paths. That made the extension boundary inconsistent and required
changes to shared modules when adding or maintaining an algorithm.
This PR refactors sparse attention around explicit module and backend contracts:
optional sparse-attention hooks for module initialization, weights, forward,
custom-op execution, and output projection.
_torch/attention_backend/sparse/<algorithm>/.SparseBackendForwardArgscarries sparse inputs from Attention/MLA to theselected backend, while
SparseRuntimeParamscarries prediction andSkipSoftmax runtime inputs from the backend to the attention op.
the hooks it needs.
module.py,backend.py,params.py,metadata.py, cache/indexer helpers,and algorithm-owned kernels. SkipSoftmax keeps its existing FMHA integration
through the shared runtime parameters.
The refactor preserves dense MLA fallback behavior, DSA cross-layer indexer
sharing and piecewise CUDA graph support, DeepSeek-V4 compressed-cache paths,
RocketKV, MiniMax-M3 integration, and SkipSoftmax.
Test Coverage
main; a native rebuild is not required forthis refactor.
pre-commiton all files changed by this PR.70 passed.
worktree.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.